# -*- coding: utf-8 -*-
"""
圆弧底形态股票筛选器
根据以下规则筛选股票：
1. 查找最近60个交易日的日线数据
2. 股票的20日线和30日线同步走出圆弧底形态
3. 当前最近的20个交易日中10日线、20日线、30日线多头排列
4. 最近30日股票尚未出现明显加速状态，最低点与最高点振幅低于15%
5. 最近30日单日最大涨幅低于8%
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager
from config import DATABASE_CONFIG
import warnings

# 忽略pandas的SettingWithCopyWarning警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'arc_uptrend_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)

class ArcBottomScreener:
    """圆弧底形态股票筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
    
    def get_stock_daily_data(self, stock_code: str, days: int = 80) -> pd.DataFrame:
        """获取股票日线数据"""
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, volume, change_pct
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换价格和数值列为float类型
            price_columns = ['open_price', 'high_price', 'low_price', 'close_price']
            for col in price_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            
            df = df.sort_values('trade_date')  # 按日期升序排列
            return df
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_ma(self, df: pd.DataFrame, periods: list = [10, 20, 30]) -> pd.DataFrame:
        """计算移动平均线"""
        if df.empty or len(df) < max(periods):
            return df
        
        df['close_price'] = pd.to_numeric(df['close_price'], errors='coerce')
        
        for period in periods:
            df[f'ma{period}'] = df['close_price'].rolling(window=period).mean()
        
        return df
    
    def is_arc_bottom(self, df: pd.DataFrame, ma_period: int = 20, check_days: int = 60) -> bool:
        """
        判断均线是否呈现圆弧底形态
        圆弧底特征：
        1. 前期下降
        2. 中间在底部盘整（找到最低点）
        3. 后期开始上升
        4. 形成U型或碗型形态
        """
        if df.empty or len(df) < check_days:
            return False
        
        ma_col = f'ma{ma_period}'
        if ma_col not in df.columns:
            return False
        
        # 获取最近check_days天的均线数据
        recent_data = df.tail(check_days)
        ma_values = recent_data[ma_col].values
        
        # 检查是否有空值
        if np.isnan(ma_values).any():
            return False
        
        # 找到最低点的位置
        min_index = np.argmin(ma_values)
        
        # 圆弧底的特征：
        # 1. 最低点不能在开头或结尾（需要有下降和上升的过程）
        # 最低点应该在中间30%-70%的位置
        if min_index < len(ma_values) * 0.3 or min_index > len(ma_values) * 0.7:
            return False
        
        # 2. 前半段（到最低点）整体应该是下降的
        first_part = ma_values[:min_index+1]
        if len(first_part) > 1:
            first_slope = (first_part[-1] - first_part[0]) / len(first_part)
            if first_slope >= 0:  # 前半段应该下降
                return False
        
        # 3. 后半段（从最低点开始）整体应该是上升的
        second_part = ma_values[min_index:]
        if len(second_part) > 1:
            second_slope = (second_part[-1] - second_part[0]) / len(second_part)
            if second_slope <= 0:  # 后半段应该上升
                return False
        
        # 4. 最近的均线值应该高于最低点（确认在上升中）
        recent_avg = np.mean(ma_values[-5:])  # 最近5天的平均值
        if recent_avg <= ma_values[min_index] * 1.01:  # 至少高于最低点1%
            return False
        
        # 5. 当前均线不能离最低点太远（避免已经大幅上涨）
        if ma_values[-1] > ma_values[min_index] * 1.15:  # 不超过最低点的15%
            return False
        
        # 6. 计算圆弧的对称性（可选）
        # 前半段和后半段的斜率绝对值应该相近
        if len(first_part) > 1 and len(second_part) > 1:
            first_abs_slope = abs(first_slope)
            second_abs_slope = abs(second_slope)
            # 斜率比例应该在合理范围内（0.5-2.0）
            if second_abs_slope > 0:
                slope_ratio = first_abs_slope / second_abs_slope
                if slope_ratio < 0.3 or slope_ratio > 3.0:
                    return False
        
        return True
    
    def is_bullish_arrangement_recent(self, df: pd.DataFrame, check_days: int = 20) -> bool:
        """判断最近N天是否为多头排列（10日线 > 20日线 > 30日线）"""
        if df.empty or len(df) < 30:
            return False
        
        # 获取最近N天的数据
        recent_data = df.tail(check_days)
        
        # 检查最近N天是否都满足多头排列
        ma_columns = ['ma10', 'ma20', 'ma30']
        
        bullish_count = 0
        total_count = 0
        
        for _, row in recent_data.iterrows():
            # 检查是否有空值
            has_nan = False
            for col in ma_columns:
                if col not in recent_data.columns or pd.isna(row[col]):
                    has_nan = True
                    break
            
            if has_nan:
                continue
            
            total_count += 1
            # 检查是否为多头排列
            if row['ma10'] > row['ma20'] > row['ma30']:
                bullish_count += 1
        
        # 至少80%的交易日满足多头排列
        if total_count == 0:
            return False
        
        return bullish_count / total_count >= 0.8
    
    def check_no_acceleration(self, df: pd.DataFrame, check_days: int = 30, 
                             max_amplitude: float = 15.0, max_single_change: float = 8.0) -> bool:
        """
        检查最近N天内是否没有明显的加速状态
        判断标准：
        1. 最近N天的最低点与最高点振幅不超过15%
        2. 最近N天单日最大涨幅不超过8%
        """
        if df.empty or len(df) < check_days:
            return False
        
        recent_data = df.tail(check_days)
        
        # 计算最近N天的最高价和最低价
        highest_price = recent_data['high_price'].max()
        lowest_price = recent_data['low_price'].min()
        
        if lowest_price == 0:
            return False
        
        # 计算振幅：(最高价 - 最低价) / 最低价 * 100
        amplitude = (highest_price - lowest_price) / lowest_price * 100
        
        # 计算最近N天的最大单日涨幅
        max_daily_change = recent_data['change_pct'].max()
        
        # 判断是否没有明显加速：
        # 1. 振幅不超过15%
        # 2. 最大单日涨幅不超过8%
        if amplitude > max_amplitude or max_daily_change > max_single_change:
            return False
        
        return True
    
    def get_stock_info(self, stock_code: str) -> dict:
        """获取股票基本信息"""
        sql = "SELECT stock_name, exchange, industry, market FROM new_stock_info WHERE stock_code = %s"
        try:
            results = self.db.execute_query(sql, (stock_code,))
            return results[0] if results else {}
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 基本信息失败: {e}")
            return {}
    
    def get_all_stock_codes(self) -> list:
        """获取所有活跃股票代码"""
        sql = "SELECT stock_code FROM new_stock_info WHERE status = 'active'"
        try:
            results = self.db.execute_query(sql)
            return [row['stock_code'] for row in results]
        except Exception as e:
            logger.error(f"获取股票代码列表失败: {e}")
            return []
    
    def screen_stocks(self) -> pd.DataFrame:
        """筛选符合条件的股票"""
        logger.info("开始筛选圆弧底形态的股票...")
        
        stock_codes = self.get_all_stock_codes()
        if not stock_codes:
            logger.error("未找到任何股票代码")
            return pd.DataFrame()
        
        logger.info(f"共找到 {len(stock_codes)} 只股票，开始筛选...")
        
        screened_stocks = []
        processed = 0
        
        for stock_code in stock_codes:
            try:
                # 获取股票数据（需要80天数据来计算均线和分析趋势）
                df = self.get_stock_daily_data(stock_code, 80)
                if df.empty or len(df) < 60:
                    continue
                
                # 计算移动平均线
                df = self.calculate_ma(df, [10, 20, 30])
                
                # 条件1：20日线圆弧底形态（检查最近60天）
                ma20_arc = self.is_arc_bottom(df, ma_period=20, check_days=60)
                if not ma20_arc:
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票，找到 {len(screened_stocks)} 只符合条件")
                    continue
                
                # 条件2：30日线圆弧底形态（检查最近60天）
                ma30_arc = self.is_arc_bottom(df, ma_period=30, check_days=60)
                if not ma30_arc:
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票，找到 {len(screened_stocks)} 只符合条件")
                    continue
                
                # 条件3：最近20个交易日中10日线、20日线、30日线多头排列
                is_bullish = self.is_bullish_arrangement_recent(df, check_days=20)
                if not is_bullish:
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票，找到 {len(screened_stocks)} 只符合条件")
                    continue
                
                # 条件4：最近30日内没有明显加速（振幅<15%，单日涨幅<8%）
                no_accel = self.check_no_acceleration(df, check_days=30, 
                                                      max_amplitude=15.0, max_single_change=8.0)
                if not no_accel:
                    processed += 1
                    if processed % 500 == 0:
                        logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票，找到 {len(screened_stocks)} 只符合条件")
                    continue
                
                # 所有条件都满足，添加到结果列表
                latest = df.iloc[-1]
                stock_info = self.get_stock_info(stock_code)
                
                # 计算额外指标
                recent_30 = df.tail(30)
                recent_60 = df.tail(60)
                
                # 计算30日振幅
                highest_30 = recent_30['high_price'].max()
                lowest_30 = recent_30['low_price'].min()
                amplitude_30 = (highest_30 - lowest_30) / lowest_30 * 100 if lowest_30 > 0 else 0
                
                # 计算30日最大单日涨幅
                max_change_30 = recent_30['change_pct'].max()
                
                # 计算均线的增长率
                ma10_values = df.tail(60)['ma10'].dropna().values
                ma20_values = df.tail(60)['ma20'].dropna().values
                ma30_values = df.tail(60)['ma30'].dropna().values
                
                ma10_slope = (ma10_values[-1] - ma10_values[0]) / ma10_values[0] * 100 if len(ma10_values) > 0 and ma10_values[0] > 0 else 0
                ma20_slope = (ma20_values[-1] - ma20_values[0]) / ma20_values[0] * 100 if len(ma20_values) > 0 and ma20_values[0] > 0 else 0
                ma30_slope = (ma30_values[-1] - ma30_values[0]) / ma30_values[0] * 100 if len(ma30_values) > 0 and ma30_values[0] > 0 else 0
                
                screened_stocks.append({
                    'stock_code': stock_code,
                    'stock_name': stock_info.get('stock_name', ''),
                    'exchange': stock_info.get('exchange', ''),
                    'industry': stock_info.get('industry', ''),
                    'market': stock_info.get('market', ''),
                    'current_price': latest['close_price'],
                    'ma10': round(latest['ma10'], 2),
                    'ma20': round(latest['ma20'], 2),
                    'ma30': round(latest['ma30'], 2),
                    'ma10_slope_60d': round(ma10_slope, 2),
                    'ma20_slope_60d': round(ma20_slope, 2),
                    'ma30_slope_60d': round(ma30_slope, 2),
                    'amplitude_30d': round(amplitude_30, 2),
                    'max_change_30d': round(max_change_30, 2),
                    'highest_30d': round(highest_30, 2),
                    'lowest_30d': round(lowest_30, 2),
                    'change_pct': latest['change_pct'],
                    'volume': latest['volume'],
                    'trade_date': latest['trade_date'].strftime('%Y-%m-%d')
                })
                
                processed += 1
                if processed % 100 == 0:
                    logger.info(f"已处理 {processed}/{len(stock_codes)} 只股票，找到 {len(screened_stocks)} 只符合条件的股票")
                    
            except Exception as e:
                logger.error(f"筛选股票 {stock_code} 失败: {e}")
                processed += 1
                continue
        
        logger.info(f"筛选完成，共找到 {len(screened_stocks)} 只符合条件的股票")
        return pd.DataFrame(screened_stocks)
    
    def save_results(self, df: pd.DataFrame, filename: str = None):
        """保存筛选结果"""
        if df.empty:
            logger.warning("没有符合条件的股票数据可保存")
            return
        
        if filename is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f'arc_uptrend_stocks_{timestamp}.csv'
        
        # 确保保存在search目录
        search_dir = os.path.dirname(os.path.abspath(__file__))
        full_path = os.path.join(search_dir, filename)
        
        try:
            df.to_csv(full_path, index=False, encoding='utf-8-sig')
            logger.info(f"筛选结果已保存到: {full_path}")
            
            # 同时保存股票代码文件
            stock_codes_filename = filename.replace('arc_uptrend_stocks_', 'stock_codes_').replace('.csv', '.txt')
            stock_codes_path = os.path.join(search_dir, stock_codes_filename)
            stock_codes = ','.join(df['stock_code'].astype(str))
            with open(stock_codes_path, 'w', encoding='utf-8') as f:
                f.write(stock_codes)
            logger.info(f"股票代码列表已保存到: {stock_codes_path}")
            
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
    
    def print_results(self, df: pd.DataFrame):
        """打印筛选结果"""
        if df.empty:
            print("未找到符合条件的股票")
            return
        
        print(f"\n=== 圆弧底形态股票筛选结果 ===")
        print(f"筛选条件：")
        print(f"  1. 20日线和30日线同步走出圆弧底形态（最近60天）")
        print(f"  2. 最近20个交易日中10日线、20日线、30日线多头排列（至少80%）")
        print(f"  3. 最近30日内振幅不超过15%")
        print(f"  4. 最近30日内单日最大涨幅不超过8%")
        print(f"共找到 {len(df)} 只符合条件的股票")
        print(f"筛选时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("\n股票列表:")
        print("-" * 200)
        print(f"{'代码':<8} {'名称':<12} {'交易所':<6} {'行业':<20} {'当前价':<8} {'MA10':<8} {'MA20':<8} {'MA30':<8} "
              f"{'MA10增长':<10} {'MA20增长':<10} {'MA30增长':<10} {'30日振幅':<10} {'30日最大涨':<10} {'日期':<12}")
        print("-" * 200)
        
        for _, row in df.iterrows():
            print(f"{row['stock_code']:<8} {row['stock_name']:<12} {row['exchange']:<6} "
                  f"{str(row['industry'])[:18]:<20} {row['current_price']:<8.2f} "
                  f"{row['ma10']:<8.2f} {row['ma20']:<8.2f} {row['ma30']:<8.2f} "
                  f"{row['ma10_slope_60d']:<10.2f} {row['ma20_slope_60d']:<10.2f} {row['ma30_slope_60d']:<10.2f} "
                  f"{row['amplitude_30d']:<10.2f} {row['max_change_30d']:<10.2f} "
                  f"{row['trade_date']:<12}")
        
        print("-" * 200)
        
        # 按行业统计
        if 'industry' in df.columns:
            industry_stats = df['industry'].value_counts()
            print(f"\n按行业分布:")
            for industry, count in industry_stats.head(10).items():
                print(f"  {industry}: {count} 只")
        
        # 按交易所统计
        if 'exchange' in df.columns:
            exchange_stats = df['exchange'].value_counts()
            print(f"\n按交易所分布:")
            for exchange, count in exchange_stats.items():
                print(f"  {exchange}: {count} 只")
        
        # 统计信息
        print(f"\n统计信息:")
        print(f"  平均股价: {df['current_price'].mean():.2f}")
        print(f"  平均MA10增长率(60天): {df['ma10_slope_60d'].mean():.2f}%")
        print(f"  平均MA20增长率(60天): {df['ma20_slope_60d'].mean():.2f}%")
        print(f"  平均MA30增长率(60天): {df['ma30_slope_60d'].mean():.2f}%")
        print(f"  平均30日振幅: {df['amplitude_30d'].mean():.2f}%")
        print(f"  平均30日最大涨幅: {df['max_change_30d'].mean():.2f}%")

def main():
    """主函数"""
    try:
        screener = ArcBottomScreener()
        
        # 筛选股票
        screened_df = screener.screen_stocks()
        
        # 打印结果
        screener.print_results(screened_df)
        
        # 保存结果
        screener.save_results(screened_df)
        
    except Exception as e:
        logger.error(f"筛选过程出错: {e}")
        raise

if __name__ == "__main__":
    main()

